import os
import ntpath
import time
import argparse
from loguru import logger
import requests
from tqdm import tqdm
from urllib.parse import quote
import concurrent.futures
from datetime import datetime


class ImageDownloader:
    def __init__(self, target, startdate, enddate, bands, localpath):
        self.target = target
        self.startdate = startdate
        self.enddate = enddate
        self.bands = bands
        self.base_localpath = localpath
        self.base_url = "http://192.168.16.70:8888/V20250420/search_by_using_obj_filter/"

    def download_file(self, file_url, target_folder, file_type=None):
        """下载指定的文件URL并保存到指定的文件夹"""
        os.makedirs(target_folder, exist_ok=True)

        filename = os.path.basename(file_url)
        if file_type == "fwhm_png":
            filename, ext = os.path.splitext(filename)
            filename = f"{filename}.fwhm{ext}"
        elif file_type == "preview_png":
            filename, ext = os.path.splitext(filename)
            filename = f"{filename}.preview{ext}"

        full_path = os.path.join(target_folder, filename)

        with requests.get(file_url, stream=True, timeout=20) as resp:
            if resp.status_code == 200:
                with open(full_path, 'wb') as f:
                    for chunk in resp.iter_content(chunk_size=8192):
                        f.write(chunk)
            else:
                logger.error(f"Download failed: {file_url} [HTTP {resp.status_code}]")
        return full_path

    def process_band(self, band):
        encoded_target = quote(self.target)
        localpath = os.path.join(self.base_localpath, self.target)
        os.makedirs(localpath, exist_ok=True)

        url = f"{self.base_url}?object={encoded_target}&filter={band}&start_date={self.startdate}&end_date={self.enddate}"
        logger.info(f"Querying [{self.target}] in band [{band}]...")

        try:
            response = requests.get(url)
            data = response.json()
        except Exception as e:
            logger.error(f"Failed to query {self.target} [{band}]: {e}")
            return

        if data.get("code") != "200" or "data" not in data:
            logger.warning(f"No results for {self.target} [{band}]")
            return

        sciimg_log_file = os.path.join(localpath, f"images_{band}.list")
        with open(sciimg_log_file, 'w') as log_file:
            log_file.write("")

        for item in tqdm(data["data"], desc=f"{self.target}-{band}"):
            try:
                obs_date = item["obs_date"]
                date_str = obs_date.replace("-", "").split(" ")[0]
                target_folder = os.path.join(localpath, date_str)

                # 下载 CalMag 文件
                cal_mag_fits = item["sciimg_sexfits"].replace(".fits", "_CalMag.fits")
                full_cal_path = os.path.join(target_folder, ntpath.basename(cal_mag_fits))
                if not os.path.exists(full_cal_path):
                    self.download_file(cal_mag_fits, target_folder)

                # 下载 SciImg 文件
                sciimg = item["sciimg"]
                full_sci_path = os.path.join(target_folder, ntpath.basename(sciimg))
                if not os.path.exists(full_sci_path):
                    self.download_file(sciimg, target_folder)

                # 写入日志
                with open(sciimg_log_file, 'a') as log_file:
                    log_file.write(full_sci_path + "\n")

            except Exception as e:
                logger.error(f"[{self.target}] processing error: {e}")

    def run(self):
        t_start = time.time()
        print(f"[{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}] 开始处理目标: {self.target}")

        with concurrent.futures.ThreadPoolExecutor() as executor:
            futures = [executor.submit(self.process_band, band) for band in self.bands]
            for future in concurrent.futures.as_completed(futures):
                try:
                    future.result()
                except Exception as e:
                    logger.error(f"Error in thread execution: {e}")

        elapsed = time.time() - t_start
        print(f"[{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}] 完成，用时：{elapsed:.2f} 秒 ≈ {elapsed / 60:.2f} 分钟")


def main():
    parser = argparse.ArgumentParser(description="下载指定天体目标的图像")
    parser.add_argument("--target", default="NGC5631", help="目标天体名称（默认: 12377）")
    parser.add_argument("--startdate", default="20221124", help="开始日期 YYYYMMDD（默认: 20221124）")
    parser.add_argument("--enddate", default="20221130", help="结束日期 YYYYMMDD（默认: 20221130）")
    parser.add_argument("--bands", default=["g", "r"], nargs="+", help="要搜索的波段（默认: g r）")
    parser.add_argument("--localpath", default="./images50cm", help="保存路径（默认: ./images50cm）")

    args = parser.parse_args()

    downloader = ImageDownloader(
        target=args.target,
        startdate=args.startdate,
        enddate=args.enddate,
        bands=args.bands,
        localpath=args.localpath
    )
    downloader.run()


if __name__ == "__main__":
    main()